package ordered_table

func countRangeSum(nums []int, lower int, upper int) int {
	sum := 0
	ans := 0
	sbt := NewBSTree()
	sbt.Add(0)
	for i := 0; i < len(nums); i++ {
		sum += nums[i]
		a := sbt.LessKeySize(sum - lower + 1)
		b := sbt.LessKeySize(sum - upper)
		ans += a - b
		sbt.Add(sum)
	}
	return ans
}

type Node struct {
	Key   int
	All   int
	Size  int
	Left  *Node
	Right *Node
}

func NewBSTree() *BSTree {
	return &BSTree{KeySet: make(map[int]struct{})}
}

type BSTree struct {
	Root   *Node
	KeySet map[int]struct{}
}

func (b *BSTree) RightRotate(cur *Node) *Node {
	if cur == nil {
		return nil
	}
	same := cur.All - b.GetLeftAll(cur) - b.GetRightAll(cur)
	leftN := cur.Left
	cur.Left = leftN.Right
	leftN.Right = cur
	leftN.Size = cur.Size
	cur.Size = b.GetLeftSize(cur) + b.GetRightSize(cur) + 1
	leftN.All = cur.All
	cur.All = b.GetLeftAll(cur) + b.GetRightAll(cur) + same
	return leftN
}

func (b *BSTree) LeftRotate(cur *Node) *Node {
	same := cur.All - b.GetLeftAll(cur) - b.GetRightAll(cur)
	rightN := cur.Right
	cur.Right = rightN.Left
	rightN.Left = cur
	rightN.Size = cur.Size
	cur.Size = b.GetLeftSize(cur) + b.GetRightSize(cur) + 1
	rightN.All = cur.All
	cur.All = b.GetLeftAll(cur) + b.GetRightAll(cur) + same
	return rightN
}

func (b *BSTree) Maintain(cur *Node) *Node {
	if cur == nil {
		return nil
	}
	leftS := b.GetLeftSize(cur)
	leftLeftS := b.GetLeftLeftSize(cur)
	leftRightS := b.GetLeftRightSize(cur)
	rightS := b.GetRightSize(cur)
	rightRightS := b.GetRightRightSize(cur)
	rightLeftS := b.GetRightLeftSize(cur)
	if rightS < leftLeftS {
		cur = b.RightRotate(cur)
		cur.Right = b.Maintain(cur.Right)
		cur = b.Maintain(cur)
	} else if rightS < leftRightS {
		cur.Left = b.LeftRotate(cur.Left)
		cur = b.RightRotate(cur)
		cur.Left = b.Maintain(cur.Left)
		cur.Right = b.Maintain(cur.Right)
		cur = b.Maintain(cur)
	} else if leftS < rightRightS {
		cur = b.LeftRotate(cur)
		cur.Left = b.Maintain(cur.Left)
		cur = b.Maintain(cur)
	} else if leftS < rightLeftS {
		cur.Right = b.RightRotate(cur.Right)
		cur = b.LeftRotate(cur)
		cur.Left = b.Maintain(cur.Left)
		cur.Right = b.Maintain(cur.Right)
		cur = b.Maintain(cur)
	}
	return cur
}

func (b *BSTree) Add(key int) {
	_, isContain := b.KeySet[key]
	b.Root = b.add(b.Root, key, isContain)
	b.KeySet[key] = struct{}{}
}

func (b *BSTree) add(cur *Node, key int, isContain bool) *Node {
	if cur == nil {
		return &Node{Key: key, All: 1, Size: 1}
	} else {
		cur.All++
		if key == cur.Key {
			return cur
		} else {
			if !isContain {
				cur.Size++
			}
			if key < cur.Key {
				cur.Left = b.add(cur.Left, key, isContain)
			} else {
				cur.Right = b.add(cur.Right, key, isContain)
			}
			return b.Maintain(cur)
		}
	}
}
func (b *BSTree) LessKeySize(key int) int {
	cur := b.Root
	ans := 0
	for cur != nil {
		if key == cur.Key {
			return ans + b.GetLeftAll(cur)
		} else if key < cur.Key {
			cur = cur.Left
		} else {
			ans += cur.All - b.GetRightAll(cur)
			cur = cur.Right
		}
	}
	return ans
}

func (b *BSTree) GetLeftSize(cur *Node) int {
	if cur.Left != nil {
		return cur.Left.Size
	}
	return 0
}

func (b *BSTree) GetLeftLeftSize(cur *Node) int {
	if cur.Left != nil && cur.Left.Left != nil {
		return cur.Left.Left.Size
	}
	return 0
}

func (b *BSTree) GetLeftRightSize(cur *Node) int {
	if cur.Left != nil && cur.Left.Right != nil {
		return cur.Left.Right.Size
	}
	return 0
}

func (b *BSTree) GetRightSize(cur *Node) int {
	if cur.Right != nil {
		return cur.Right.Size
	}
	return 0
}

func (b *BSTree) GetRightRightSize(cur *Node) int {
	if cur.Right != nil && cur.Right.Right != nil {
		return cur.Right.Right.Size
	}
	return 0
}

func (b *BSTree) GetRightLeftSize(cur *Node) int {
	if cur.Right != nil && cur.Right.Left != nil {
		return cur.Right.Left.Size
	}
	return 0
}

func (b *BSTree) GetLeftAll(cur *Node) int {
	if cur.Left != nil {
		return cur.Left.All
	}
	return 0
}

func (b *BSTree) GetRightAll(cur *Node) int {
	if cur.Right != nil {
		return cur.Right.All
	}
	return 0
}

// ################## extend for medianSlidingWindow   ################

func (b *BSTree) GetIndexKey(index int) float64 {
	cur := b.Root
	for cur != nil {
		curSame := cur.All - b.GetLeftAll(cur) - b.GetRightAll(cur)
		if cur.Left != nil && cur.Left.All < index && cur.Left.All + curSame >= index || (cur.Left == nil && curSame >= index) {
			return float64(cur.Key)
		} else if cur.Left != nil && cur.Left.All+1 > index {
			cur = cur.Left
		} else {
			same := cur.All - b.GetLeftAll(cur) - b.GetRightAll(cur)
			index = index - same - b.GetLeftAll(cur)
			cur = cur.Right
		}
	}
	return 0.0
}

func (b *BSTree) Remove(key int) {
	_, contain := b.KeySet[key]
	if contain {
		var isRemoveNode bool
		b.Root, isRemoveNode = b.delete(b.Root, key)
		if isRemoveNode && b.Root != nil {
			b.ReCalcSize(b.Root)
		}
	}
}

func (b *BSTree) delete(cur *Node, key int) (*Node, bool) {
	if cur == nil {return nil, false}
	var isRemoveNode bool
	cur.All--
	if key < cur.Key {
		cur.Left, isRemoveNode = b.delete(cur.Left, key)
		if isRemoveNode && cur.Left != nil {
			b.ReCalcSize(cur.Left)
		}
	} else if key > cur.Key {
		cur.Right, isRemoveNode = b.delete(cur.Right, key)
		if isRemoveNode && cur.Right != nil {
			b.ReCalcSize(cur.Right)
		}
	} else {
		curSame := cur.All - b.GetLeftAll(cur) - b.GetRightAll(cur)
		if curSame == 0 {
			isRemoveNode = true
			if cur.Left == nil && cur.Right == nil {
				cur = nil
			} else if cur.Left != nil && cur.Right == nil {
				cur = cur.Left
			} else if cur.Left == nil && cur.Right != nil {
				cur = cur.Right
			} else {
				var pre *Node
				des := cur.Right
				des.Size--
				for des.Left != nil {
					pre = des
					des = des.Left
					des.Size--
				}
				desSame := des.All - b.GetRightAll(des)
				des2 := cur.Right
				des2.All -= desSame
				for des2.Left != nil {
					des2 = des2.Left
					des2.All -= desSame
				}
				if pre != nil {
					pre.Left = des.Right
					des.Right = cur.Right
				}
				des.Left = cur.Left
				des.Size = des.Left.Size + b.GetRightSize(des) + 1
				des.All = b.GetRightAll(des) + b.GetLeftAll(des) + desSame
				cur = des
			}
		}
	}
	return cur, isRemoveNode
}

func (b *BSTree) ReCalcSize(cur *Node) {
	if cur == nil {
		return
	}
	cur.Size = b.GetRightSize(cur) + b.GetLeftSize(cur) + 1
}
